# 加载环境变量
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

from Agent.ReAct import ReActAgent
from Models.Factory import ChatModelFactory
from Tools import *
from Tools.BankCustomerMockTool import BankCustomerMockInentTool
from Tools.RagTool import RagTool
from langchain_community.chat_message_histories.in_memory import ChatMessageHistory


def launch_agent(agent: ReActAgent, welcome_message: str = "有什么可以帮您？"):
    human_icon = "\U0001F468"
    ai_icon = "\U0001F916"
    chat_history = ChatMessageHistory()

    first_turn = True
    while True:
        if not first_turn:
            task = input(f"{human_icon}：")
        else:
            task = input(f"{ai_icon}：{welcome_message}\n{human_icon}：")
            first_turn = False
            
        if task.strip() == "":
            continue
        
        if task.strip().lower() == "quit":
            break
        reply = agent.run(task, chat_history, verbose=True)
        print(f"{ai_icon}：{reply}\n")
        
def launch_agents(assistant_agent: ReActAgent, customer_agent: ReActAgent, welcome_message: str = "有什么可以帮您？"):
    human_icon = "\U0001F468"
    ai_icon = "\U0001F916"
    chat_history = ChatMessageHistory()

    turns = []
    print(f"{ai_icon}：{welcome_message}")
    turns.append(f"{ai_icon}：{welcome_message}")
    reply = welcome_message
    while True:
        print(f"\033[31m [用户 agent]] \033[0m")
        request = customer_agent.run(reply, chat_history, verbose=False)
        print(f"{human_icon}：{request}\n")
        turns.append(f"{human_icon}：{request}\n")
        
        if request.find("没有其他需要") >= 0:
            break
        
        print(f"\033[31m [客服 agent]] \033[0m")
        reply = assistant_agent.run(request, chat_history, verbose=False)
        print(f"{ai_icon}：{reply}\n")
        turns.append(f"{ai_icon}：{reply}\n")
        
        if reply.find("祝您生活愉快") >= 0:
            break
    
    print("\n" + "\n".join(turns))

def main():
    # 语言模型
    llm = ChatModelFactory.get_model("qwen-max")

    # 自定义工具集
    tools = [
        verify_code_tool,
        query_saving_account_by_id_tool,
        query_credit_account_by_id_tool,
        document_qa_tool,
        finish_placeholder,
        request_placeholder,
        document_qa_tool,
        directory_inspection_tool
    ]

    # 定义智能体
    agent = ReActAgent(
        llm=llm,
        tools=tools,
        work_dir="./data",
        main_prompt_file="./prompts/bank_assistant.txt",
        max_thought_steps=10,
    )

    # 运行智能体
    launch_agent(agent)

def main2():
    # 语言模型
    llm = ChatModelFactory.get_model("qwen-max")

    # assistant tools
    assistant_tools = [
        verify_code_tool,
        query_saving_account_by_id_tool,
        query_credit_account_by_id_tool,
        document_qa_tool,
        finish_placeholder,
        request_placeholder,
        document_qa_tool,
        directory_inspection_tool
    ]

    # assitant agent
    assistant_agent = ReActAgent(
        llm=llm,
        tools=assistant_tools,
        work_dir="./bankdata",
        main_prompt_file="./prompts/bank_assistant.txt",
        max_thought_steps=10
    )
    
    # customer tools
    customer_tools = [
        provide_my_info_tool,
        provide_my_verify_code_tool,
        BankCustomerMockInentTool().as_tool(),
        provide_process_account_type_tool,
        finish_placeholder
    ]
    
    # customer agent
    customer_agent = ReActAgent(
        llm=llm,
        tools=customer_tools,
        work_dir="./bankdata",
        main_prompt_file="./prompts/bank_customer.txt",
        max_thought_steps=3,
    )

    # 运行智能体
    launch_agents(assistant_agent, customer_agent, "您好，橘猫银行客服，请问有什么可以帮您?")
    #launch_agent(assistant_agent)
    
if __name__ == "__main__":
    main()
